import torch
import torch.nn as nn

from typing import Tuple, Dict
from . import model_utils
from .model_opts import ModelOpts


def _gen_rigid_features() -> torch.Tensor:
    # predefined rigid samples on axes and planes
    samples = torch.tensor([
        [0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0],
        [0.0, 0.0, 1.0],
        [-1.0, -1.0, 0.0],
        [-1.0, 0.0, -1.0],
        [0.0, -1.0, -1.0],
    ], dtype=torch.float)
    # scale to expected norm 0.1
    norms = samples.norm(dim=1, keepdim=True)
    scale = torch.where(norms > 0, 0.1 / norms, torch.zeros_like(norms))
    return samples * scale


def get_output_index_ranges(
        mo: ModelOpts,
        predict_skel_scale: bool,
) -> Tuple[Dict[str, Tuple[int, int]], int]:
    rigid_samples = _gen_rigid_features()

    assert mo.nWristRigidPts <= len(rigid_samples), (
        "Max supported n_wrist_rigid_pts is "
        f"{len(rigid_samples)}, got {mo.nWristRigidPts}"
    )

    output_dims = {
        "joint_angles": 20,
        "wrist_xfs": mo.nWristRigidPts * 3,
        "skel_scales": 1 if predict_skel_scale else 0,
        # "landmark_uncertainty_sigmas": 21,
        "hand_id": 2
    }
    n_output_dims = 0
    output_index_range = {}
    for k, v in output_dims.items():
        if v != 0:
            output_index_range[k] = (n_output_dims, n_output_dims + v)
            n_output_dims = n_output_dims + v
    return output_index_range, n_output_dims


def get_output_index_ranges_sk(
        mo: ModelOpts,
        predict_skel_scale: bool,
) -> Tuple[Dict[str, Tuple[int, int]], int]:
    rigid_samples = _gen_rigid_features()

    assert mo.nWristRigidPts <= len(rigid_samples), (
        "Max supported n_wrist_rigid_pts is "
        f"{len(rigid_samples)}, got {mo.nWristRigidPts}"
    )

    output_dims = {
        "joint_angles": 20,
        "wrist_xfs": mo.nWristRigidPts * 3,
        "skel_scales": 1 if predict_skel_scale else 0,
        "known_hand_model_num": mo.nKnownHandModels,
    }
    n_output_dims = 0
    output_index_range = {}
    for k, v in output_dims.items():
        if v != 0:
            output_index_range[k] = (n_output_dims, n_output_dims + v)
            n_output_dims = n_output_dims + v
    return output_index_range, n_output_dims


class PoseRegressor(nn.Module):
    def __init__(
            self,
            mo: ModelOpts,
            predict_skel_scale: bool,
    ):
        super().__init__()
        # prepare output index ranges
        ranges, n_out = get_output_index_ranges(mo, predict_skel_scale)
        self.register_buffer("_left_wrist_samples", _gen_rigid_features()[: mo.nWristRigidPts])
        self._ranges = ranges
        # regression head
        self.n_in_channel = mo.nImageFeatureChannels if predict_skel_scale else mo.nImageFeatureChannels + mo.nSkeletonFeatureChannels
        self._head = model_utils.create_pose_regression_layers(
            n_in_channels=self.n_in_channel,
            n_blocks=mo.nPoseRegressionBlocks,
            n_out_channels=n_out,
        )
        self._predict_skel_scale = predict_skel_scale
        self._softmax = nn.Softmax(dim=1)

    def forward(self, img_skel_feats: torch.Tensor):
        # img_skel_feats: [B, C, H, W]
        B = img_skel_feats.size(0)
        feats = self._head(img_skel_feats)
        feats = feats.flatten(start_dim=1)

        # joint angles
        ja_start, ja_end = self._ranges["joint_angles"]
        finger = feats[:, ja_start:ja_end]
        wrist_zero = torch.zeros(B, 2, device=feats.device, dtype=feats.dtype)
        joint_angles = torch.cat([finger, wrist_zero], dim=1)

        # wrist xfs via SVD Procrustes
        wx_start, wx_end = self._ranges["wrist_xfs"]
        raw_wrist = feats[:, wx_start:wx_end]
        rigid_pts = raw_wrist.view(B, -1, 3)
        # expand source
        src = self._left_wrist_samples.unsqueeze(0).expand(B, -1, -1)
        wrist_xfs = model_utils.procrustes_align(src, rigid_pts)

        skel_scales = None
        if self._predict_skel_scale:
            ss_start, ss_end = self._ranges["skel_scales"]
            log_sc = feats[:, ss_start:ss_end].squeeze(1)
            skel_scales = torch.exp(log_sc)

        # hand_id
        h_start, h_end = self._ranges["hand_id"]
        hand_side_logits = feats[:, h_start:h_end]
        # hand_side_probs = self._softmax(hand_side_logits)

        # return joint_angles, wrist_xfs, skel_scales, sigmas
        return joint_angles, wrist_xfs, skel_scales, hand_side_logits


class PoseRegressorDeploy(nn.Module):
    def __init__(
            self,
            mo: ModelOpts,
            predict_skel_scale: bool,
    ):
        super().__init__()
        # prepare output index ranges
        ranges, n_out = get_output_index_ranges(mo, predict_skel_scale)
        self.register_buffer("_left_wrist_samples", _gen_rigid_features()[: mo.nWristRigidPts])
        self._ranges = ranges
        # regression head
        self._head = model_utils.create_pose_regression_layers(
            n_in_channels=mo.nImageFeatureChannels,
            n_blocks=mo.nPoseRegressionBlocks,
            n_out_channels=n_out,
        )
        self._predict_skel_scale = predict_skel_scale
        self._softmax = nn.Softmax(dim=1)

    def forward(self, img_skel_feats: torch.Tensor):
        # img_skel_feats: [B, C, H, W]
        B = img_skel_feats.size(0)
        feats = self._head(img_skel_feats)
        feats = feats.flatten(start_dim=1)

        # joint angles
        ja_start, ja_end = self._ranges["joint_angles"]
        finger = feats[:, ja_start:ja_end]
        wrist_zero = torch.zeros(B, 2, device=feats.device, dtype=feats.dtype)
        joint_angles = torch.cat([finger, wrist_zero], dim=1)

        # wrist xfs via SVD Procrustes
        wx_start, wx_end = self._ranges["wrist_xfs"]
        raw_wrist = feats[:, wx_start:wx_end]
        rigid_pts = raw_wrist.view(B, -1, 3)

        skel_scales = None
        if self._predict_skel_scale:
            ss_start, ss_end = self._ranges["skel_scales"]
            log_sc = feats[:, ss_start:ss_end].squeeze(1)
            skel_scales = torch.exp(log_sc)

        # hand_id
        h_start, h_end = self._ranges["hand_id"]
        hand_side_logits = feats[:, h_start:h_end]
        # hand_side_probs = self._softmax(hand_side_logits)

        # return joint_angles, wrist_xfs, skel_scales, sigmas
        return joint_angles, rigid_pts, skel_scales, hand_side_logits


class PoseRegressorSK(nn.Module):
    def __init__(
            self,
            mo: ModelOpts,
            predict_skel_scale: bool,
    ):
        super().__init__()
        # prepare output index ranges
        ranges, n_out = get_output_index_ranges_sk(mo, predict_skel_scale)
        self.register_buffer("_left_wrist_samples", _gen_rigid_features()[: mo.nWristRigidPts])
        self._ranges = ranges
        # regression head
        self.n_in_channel = mo.nImageFeatureChannels + mo.nSkeletonFeatureChannels
        self._head = model_utils.create_pose_regression_layers(
            n_in_channels=self.n_in_channel,
            n_blocks=mo.nPoseRegressionBlocks,
            n_out_channels=n_out,
        )
        self._predict_skel_scale = predict_skel_scale
        self._softmax = nn.Softmax(dim=1)

    def forward(self, img_skel_feats: torch.Tensor):
        # img_skel_feats: [B, C, H, W]
        B = img_skel_feats.size(0)
        feats = self._head(img_skel_feats)
        feats = feats.flatten(start_dim=1)

        # joint angles
        ja_start, ja_end = self._ranges["joint_angles"]
        finger = feats[:, ja_start:ja_end]
        wrist_zero = torch.zeros(B, 2, device=feats.device, dtype=feats.dtype)
        joint_angles = torch.cat([finger, wrist_zero], dim=1)

        # wrist xfs via SVD Procrustes
        wx_start, wx_end = self._ranges["wrist_xfs"]
        raw_wrist = feats[:, wx_start:wx_end]
        rigid_pts = raw_wrist.view(B, -1, 3)
        # expand source
        src = self._left_wrist_samples.unsqueeze(0).expand(B, -1, -1)
        wrist_xfs = model_utils.procrustes_align(src, rigid_pts)

        skel_scales = None
        if self._predict_skel_scale:
            ss_start, ss_end = self._ranges["skel_scales"]
            log_sc = feats[:, ss_start:ss_end].squeeze(1)
            skel_scales = torch.exp(log_sc)

        # hand_id
        h_start, h_end = self._ranges["known_hand_model_num"]
        hand_model_logits = feats[:, h_start:h_end]
        # hand_model_probs = self._softmax(hand_model_logits)

        # return joint_angles, wrist_xfs, skel_scales, sigmas
        return joint_angles, wrist_xfs, skel_scales, hand_model_logits
